import matplotlib.pyplot as plt
import pandas as pd
import sqlite3
import matplotlib.dates as mdates  # 导入日期格式化模块
import akshare as ak

def fund_flow(ts_name = "振华科技",ts_code = "000733"):
    plt.rcParams["font.sans-serif"] = ["Songti SC"]  # 设置中文字体（Windows）
    plt.rcParams["axes.unicode_minus"] = False    # 解决负号显示异常
    # 首先判断数据库存不存在
    with sqlite3.connect("example.db") as conn:
        # 不存在的话使用查询akshare
        ifExists = pd.read_sql(f"select name from sqlite_master where type = 'table' and name = 'df_stock_{ts_code}'",conn)
        if ifExists.empty:
            df_stock = ak.stock_individual_fund_flow(stock=ts_code, market= "sz")
            # print(df_stock.head())
            df_stock.to_sql(f"df_stock_{ts_code}",conn,if_exists="replace",index=False)
        df_stock = pd.read_sql(f"select * from df_stock_{ts_code}",conn)
        
    # 转换日期格式（如果日期是字符串）
    df_stock["日期"] = pd.to_datetime(df_stock["日期"])
    # 设置画布大小
    plt.figure(figsize=(48, 6))
    # 绘制折线
    plt.plot(df_stock["日期"], df_stock["主力净流入-净额"], label="主力资金", marker="o", color="#2C73D2")
    plt.plot(df_stock["日期"], df_stock["大单净流入-净额"], label="大单资金", linestyle="--", color="#4B4453")
    plt.plot(df_stock["日期"], df_stock["中单净流入-净额"], label="中单资金", marker="s", color="#FF8066")
    plt.plot(df_stock["日期"], df_stock["小单净流入-净额"], label="小单资金", linestyle="-.", color="#845EC2")
    
    # --- 关键修改：设置 x 轴日期格式 ---
    # 定义日期格式（年-月-日）
    date_format = mdates.DateFormatter("%Y-%m-%d")
    plt.gca().xaxis.set_major_formatter(date_format)
    
    # 自动调整 x 轴刻度密度（避免日期重叠）
    # plt.gca().xaxis.set_major_locator(mdates.AutoDateLocator())
    # 在折线图中设置每隔 N 天显示一个标签
    plt.gca().xaxis.set_major_locator(mdates.DayLocator(interval=1))  # 每3天显示一个日期
    # 添加标题和标签
    plt.title(f"{ts_name}({ts_code}) 资金流动趋势", fontsize=14, fontweight="bold")
    plt.xlabel("日期", fontsize=5)
    plt.ylabel("净流入金额 (万元)", fontsize=12)
    plt.xticks(rotation=90)  # 旋转日期标签避免重叠
    
    # 添加图例和网格
    plt.legend(loc="upper left", fontsize=10)
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    
    # 自动调整布局
    plt.tight_layout()
    #plt.show()
    plt.savefig(f"{ts_code}_plot.svg",format='svg')
    
    # 假设已获取数据 df_last5（包含日期列）
    df_last5 = df_stock.tail(5)
    
    # 将日期转为字符串（示例：'2023-01-01'）
    dates = df_last5["日期"].dt.strftime("%Y-%m-%d").tolist()
    
    # 绘制柱状图时，直接使用日期字符串作为 x 轴标签
    x = range(len(dates))  # x 位置仍用数字坐标
    width = 0.2
    
    plt.figure(figsize=(12, 6))
    plt.bar(x, df_last5["主力净流入-净额"], width, label="主力资金", color="#2C73D2")
    plt.bar([i + width for i in x], df_last5["大单净流入-净额"], width, label="大单资金", color="#4B4453")
    plt.bar([i + 2*width for i in x], df_last5["中单净流入-净额"], width, label="中单资金", color="#FF8066")
    plt.bar([i + 3*width for i in x], df_last5["小单净流入-净额"], width, label="小单资金", color="#845EC2")
    
    # 设置 x 轴标签为日期字符串
    plt.xticks([i + 1.5*width for i in x], dates, rotation=45)  # 直接传递 dates 列表
    
    # 其他设置保持不变
    plt.title(f"{ts_name}({ts_code}) 近5个交易日资金净流入对比", fontsize=14, fontweight="bold")
    plt.xlabel("日期", fontsize=12)
    plt.ylabel("净流入金额 (万元)", fontsize=12)
    plt.legend(loc="upper right", fontsize=10)
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    plt.tight_layout()
    #plt.show()
    plt.savefig(f"{ts_code}_bar.svg",format='svg')
    
    # 提取最新一日数据
    latest_data = df_stock.iloc[-1]
    
    # 定义标签和金额（注意：此处保留符号）
    labels = ["主力资金", "大单资金", "中单资金", "小单资金"]
    sizes = [
        latest_data["主力净流入-净额"],
        latest_data["大单净流入-净额"],
        latest_data["中单净流入-净额"],
        latest_data["小单净流入-净额"]
    ]
    
    # 过滤掉净流入为0的项，并取绝对值
    abs_sizes = [abs(size) for size in sizes]
    labels = [label for label, size in zip(labels, sizes) if size != 0]
    abs_sizes = [abs(size) for size in sizes if size != 0]
    
    # 根据原始数值的符号设置颜色（流入绿色，流出红色）
    colors = ["#4CAF50" if size > 0 else "#F44336" for size in sizes if size != 0]
    
    # 绘制饼图
    plt.figure(figsize=(8, 8))
    wedges, texts, autotexts = plt.pie(
        abs_sizes,
        labels=labels,
        autopct="%1.1f%%",
        startangle=90,
        colors=colors,
        shadow=True,
        textprops={"fontsize": 12}
    )
    
    # 添加图例说明颜色含义
    plt.legend(
        handles=[
            plt.Rectangle((0,0),1,1, fc="#4CAF50", alpha=0.5), 
            plt.Rectangle((0,0),1,1, fc="#F44336", alpha=0.5)
        ],
        labels=["资金流入", "资金流出"],
        loc="upper left"
    )
    
    plt.title(f"{latest_data['日期'].strftime('%Y-%m-%d')} 资金流动强度（绝对值占比）", fontsize=14, fontweight="bold")
    # plt.show()
    plt.savefig(f"{ts_code}_pie.svg",format='svg')
    ...


if __name__ == '__main__':
    fund_flow()
